package edu.bu.manateem.msg.heatmap;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.imageio.ImageIO;

import org.apache.batik.dom.svg.SVGDOMImplementation;
import org.apache.batik.dom.util.DOMUtilities;
import org.apache.commons.collections.primitives.ArrayFloatList;
import org.apache.commons.collections.primitives.ArrayIntList;
import org.systemsbiology.jrap.MSXMLParser;
import org.w3c.dom.DOMImplementation;
import org.w3c.dom.Element;
import org.w3c.dom.svg.SVGDocument;

import edu.bu.manateem.msg.model.Compound;
import edu.bu.manateem.msg.model.MassScan;

public class HeatMapBuilder
	//extends Thread
{
	private ArrayIntList scanList = new ArrayIntList();
	private ArrayIntList mzList = new ArrayIntList();
	private ArrayFloatList intensityList = new ArrayFloatList();
	
	private Map<Integer,MassScan> scans = new HashMap<Integer,MassScan>();
	
	private Set<HeatmapTarget> heatMapTargets = new HashSet<HeatmapTarget>();
	
	private String mzFileName;
	private String saveFileName;
	
	private double maxMz = 1200;
	private double binSize = 0.05;
	private Integer maxBin = null;
	
	Integer minScanToPlot = null;
	Integer maxScanToPlot = null;
	
	private Integer heatmapWidth = null;
	private Integer heatmapHeight = null;
	
	public HeatMapBuilder(String mzFileName,String saveFileName)
	{
		this.mzFileName = mzFileName;
		this.saveFileName = saveFileName;
	}
	
	// public static void main(String[] args)
	// {
	// String fileName =
	// "/Users/gindiny/Documents/zaia/experimental/Aorta_control_2.mzXML";
	// HeatMapBuilder builder = new HeatMapBuilder(fileName);
	// builder.setMinScanToPlot(6000);
	// builder.setMaxScanToPlot(6600);
	// builder.setBinSize(0.1);
	// builder.start();
	// }
	
	/**
	 * From msInspect
	 * @param x
	 * @param fABS
	 * @return
	 */
	private static float medianSampled(ArrayFloatList x,boolean fABS)
	{
		double p = 0.5;
		int len = x.size();
		int sample = 100 + (int) Math.ceil(Math.sqrt(len));
		
		float[] copy = null;
		if (sample < len / 4)
		{
			copy = new float[sample];
			for (int i = 0; i < sample; i++)
			{
				int r = (int) Math.floor(Math.random() * len);
				copy[i] = fABS
					? Math.abs(x.get(r))
					: x.get(r);
			}
		}
		else
		{
			copy = new float[len];
			for (int i = 0; i < len; i++)
			{
				copy[i] = fABS
					? Math.abs(x.get(i))
					: x.get(i);
			}
		}
		Arrays.sort(copy);
		double l = (copy.length * p) - 1;
		int i = Math.max(0,(int) Math.floor(l));
		return copy[i];
	}
	
	public void addHeatMapTarget(Compound species,int scan,int numberOfScans)
	{
		int mZ = getBin((float) species.getMz());
		HeatmapTarget target =
			new HeatmapTarget(scan,mZ,numberOfScans,species);
		heatMapTargets.add(target);
	}
	
	public void addScan(Integer index,MassScan scan)
	{
		scans.put(index,scan);
	}
	
	private void generateSvgHeatmap(File imageFile)
	{
		DOMImplementation impl = SVGDOMImplementation.getDOMImplementation();
		String svgNS = SVGDOMImplementation.SVG_NAMESPACE_URI;
		SVGDocument doc = (SVGDocument) impl.createDocument(svgNS,"svg",null);
		
		Element rootElement = doc.getDocumentElement();
		rootElement.setAttributeNS(null,"width",Integer.toString(heatmapWidth));
		rootElement.setAttributeNS(null,"height",
			Integer.toString(heatmapHeight));
		
		// Link to heatmap
		Element imageElement =
			doc.createElementNS("http://www.w3.org/1999/xlink","image");
		imageElement.setAttributeNS(null,"xlink:href",imageFile.getName());
		imageElement.setAttributeNS(null,"x","0");
		imageElement.setAttributeNS(null,"y","0");
		imageElement.setAttributeNS(null,"height",
			Integer.toString(heatmapHeight));
		imageElement.setAttributeNS(null,"width",Integer.toString(heatmapWidth));
		rootElement.appendChild(imageElement);
		
		for (HeatmapTarget target : heatMapTargets)
		{
			int oldX = target.getxCoordinate();
			target.setxCoordinate(oldX - minScanToPlot);
			rootElement.appendChild(target.getSvgElement(doc,heatmapHeight));
		}
		// trying out scale
		HeatmapScale scale = new HeatmapScale(doc,heatmapHeight,10,binSize);
		rootElement.appendChild(scale.getVerticalLine());
		
		// ticks
		List<Element> ticks = scale.getTicks();
		for (Element tick : ticks)
		{
			rootElement.appendChild(tick);
		}
		
		List<Element> lables = scale.getLabels();
		
		for (Element label : lables)
		{
			rootElement.appendChild(label);
		}
		
		try
		{
			PrintWriter printWriter =
				new PrintWriter(new FileWriter(saveFileName + ".svg"));
			DOMUtilities.writeDocument(doc,printWriter);
			printWriter.close();
		}
		catch (IOException e)
		{
			e.printStackTrace();
		}
		
	}
	
	private int getBin(float value)
	{
		int bin = (int) Math.ceil(value / binSize);
		return Math.min(bin,getMaxBin());
	}
	
	/**
	 * @return the binSize
	 */
	public double getBinSize()
	{
		return binSize;
	}
	
	private File getImageFile()
	{
		float medianSampled = HeatMapBuilder.medianSampled(intensityList,false);
		float threshold = medianSampled / 2.0f;
		
		BufferedImage heatmapImage =
			ReducedIntensityPlot.generateImage(threshold,scanList,mzList,
				intensityList);
		heatmapHeight = heatmapImage.getHeight();
		heatmapWidth = heatmapImage.getWidth();
		
		File imageFile = new File(saveFileName + ".png");
		
		try
		{
			ImageIO.write(heatmapImage,"png",imageFile);
		}
		catch (IOException e)
		{
			e.printStackTrace();
		}
		
		return imageFile;
	}
	
	/**
	 * @return the maxBin
	 */
	public int getMaxBin()
	{
		if (maxBin == null)
		{
			maxBin = (int) Math.ceil(maxMz / binSize);
		}
		
		return maxBin;
	}
	
	/**
	 * @return the maxMz
	 */
	public double getMaxMz()
	{
		return maxMz;
	}
	
	/**
	 * @return the maxScanToPlot
	 */
	public Integer getMaxScanToPlot()
	{
		return maxScanToPlot;
	}
	
	/**
	 * @return the minScanToPlot
	 */
	public Integer getMinScanToPlot()
	{
		return minScanToPlot;
	}
	
	/**
	 * Scans the mzXML file data and builds a data
	 * representation
	 */
	private void initImage()
	{
		MSXMLParser msXMLparser = new MSXMLParser(mzFileName);
		
		if (minScanToPlot == null)
		{
			minScanToPlot = 1;
		}
		if (maxScanToPlot == null)
		{
			maxScanToPlot = msXMLparser.getScanCount();
		}
		
		MassScan scan;
		for (int scanIndex = minScanToPlot,loopIndex = 1; scanIndex <= maxScanToPlot; scanIndex++,loopIndex++)
		{
			if (scans.containsKey(scanIndex))
			{
				scan = scans.get(scanIndex);
			}
			else
			{
				scan = new MassScan(msXMLparser.rap(scanIndex));
			}
			
			float[][] spectrum = scan.getScanValues();
			
			plotColumn(loopIndex - 1,spectrum);
		}
	}
	
	private void plotColumn(int scan,float[][] spectrum)
	{
		int length = spectrum[0].length;
		if (length == 0)
		{
			return;
		}
		
		Integer mzCurrent = (int) spectrum[0][0];
		float intensityMax = (int) spectrum[1][0];
		
		for (int i = 1; i < length; i++)
		{
			int mz = getBin(spectrum[0][i]);
			float intensity = spectrum[1][i];
			
			if (mz == mzCurrent)
			{
				intensityMax = intensityMax > intensity
					? intensityMax
					: intensity;
			}
			else
			{
				if (intensityMax > 1)
				{
					scanList.add(scan);
					mzList.add(mzCurrent);
					intensityList.add(intensityMax);
				}
				mzCurrent = mz;
				intensityMax = intensity;
			}
		}
		scanList.add(scan);
		mzList.add(mzCurrent);
		intensityList.add(intensityMax);
	}
	
//	@Override
	public void run()
	{
		initImage();
		File imageFile = getImageFile();
		generateSvgHeatmap(imageFile);
	}
	
	/**
	 * @param binSize the binSize to set
	 */
	public void setBinSize(double binSize)
	{
		this.binSize = binSize;
	}
	
	/**
	 * @param maxBin the maxBin to set
	 */
	public void setMaxBin(int maxBin)
	{
		this.maxBin = maxBin;
	}
	
	/**
	 * @param maxMz the maxMz to set
	 */
	public void setMaxMz(double maxMz)
	{
		this.maxMz = maxMz;
	}
	
	/**
	 * @param maxScanToPlot the maxScanToPlot to set
	 */
	public void setMaxScanToPlot(Integer maxScanToPlot)
	{
		this.maxScanToPlot = maxScanToPlot;
	}
	
	/**
	 * @param minScanToPlot the minScanToPlot to set
	 */
	public void setMinScanToPlot(Integer minScanToPlot)
	{
		this.minScanToPlot = minScanToPlot;
	}
}
